package com.bins.langchain.langchain4j.config;

import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class MilvusConfig {

    private final MilvusProperties milvusProperties;

    public MilvusConfig(MilvusProperties milvusProperties) {
        this.milvusProperties = milvusProperties;
    }

    @Bean
    public MilvusServiceClient milvusServiceClient() {
        return new MilvusServiceClient(
                ConnectParam.newBuilder()
                        .withHost(milvusProperties.getHost())
                        .withPort(milvusProperties.getPort())
                        .withDatabaseName(milvusProperties.getDatabaseName())
                        .withAuthorization(milvusProperties.getUsername(), milvusProperties.getPassword())
                        .build());
    }

    @Bean
    public MilvusEmbeddingStore milvusEmbeddingStore() {
        MilvusEmbeddingStore store = MilvusEmbeddingStore.builder()
                .milvusClient(milvusServiceClient())
                .collectionName(milvusProperties.getCollectionName())
                .dimension(milvusProperties.getDimension())
                .indexType(IndexType.IVF_FLAT)
                .build();
        return store;
    }

}
